"""
write by qianqianjun
使用预训练的 VGG 网络 作为 DAE的编码器实现
"""

import torch
import torch.nn as nn
from autoencoder.parameter import Parameter
from autoencoder.parameter import parameter as parm
from collections import OrderedDict
from autoencoder.utils import GridSpatialIntegral, Warper,Mixer


####### VGG 编码器体系结构 ########
class Encoders(nn.Module):
    def __init__(self,parm:Parameter):
        """
        DAE 的编码器
        :param parm: 超参数集合
        """
        super(Encoders, self).__init__()
        self.vgg=VggEncoder(parm,True)
        self.flattener=VggFlattener(parm)
        self.Imixer=Mixer(parm,parm.zdim,parm.idim)
        self.Wmixer=Mixer(parm,parm.zdim,parm.wdim)

    def forward(self, x):
        conv3_1,conv4_1,conv5_1,out= self.vgg(x)
        self.z=self.flattener(out)
        self.zImg=self.Imixer(self.z)
        self.zWarp=self.Wmixer(self.z)
        return self.z,self.zImg,self.zWarp

class VggEncoder(nn.Module):
    def __init__(self,parm:Parameter,load_pretrain:bool=True):
        super(VggEncoder, self).__init__()
        self.parm=parm
        self.block1=self.block(name=1,layer=2,channel=64,in_channel=3)
        self.block2=self.block(name=2,layer=2,channel=128,in_channel=64)
        self.conv3_1=self.deep(128,256,3)
        self.block3=self.block(name=3,layer=3,channel=256,in_channel=256,start=2)
        self.conv4_1=self.deep(256,512,4)
        self.block4=self.block(name=4,layer=3,channel=512,in_channel=512,start=2)
        self.conv5_1=self.deep(512,512,5)
        self.block5=self.block(name=5,layer=3,channel=512,in_channel=512,start=2)

        if load_pretrain:
            self._load_pretrain()

    def _load_pretrain(self):
        load_dict = torch.load(parm.pretrainedModel)
        state_dic = self.state_dict()
        step = 0
        for i, j in zip(state_dic, load_dict):
            step += 1
            if step == 32:
                break
            state_dic[i] = load_dict[j]
        self.load_state_dict(state_dic)
        print("加载完成")

    def block(self,name:int,layer:int,channel:int,in_channel,start=1):
        layers=[]
        for l in range(start,layer+start):
            if l!=1:
                in_channel=channel
            layers.append(("conv{}_{}".format(name,l),nn.Conv2d(
                in_channels=in_channel,out_channels=channel,kernel_size=3,
                stride=1,padding=1
            )))
            layers.append(("norm{}_{}".format(name,l),nn.InstanceNorm2d(channel)))
            layers.append(("relu{}_{}".format(name,l),nn.LeakyReLU(0.2,True)))

        layers.append(("maxpool{}".format(name), nn.MaxPool2d(2, 2)))

        return nn.Sequential(OrderedDict(layers))
    def deep(self,in_channels,out_channels,name):
        return nn.Sequential(OrderedDict([
            ("conv{}_1".format(name),nn.Conv2d(in_channels,out_channels,3,1,1)),
            ("norm{}_1".format(name),nn.InstanceNorm2d(out_channels)),
            ("relu{}_1".format(name),nn.LeakyReLU(0.2,True))
        ]))

    def forward(self, x:torch.Tensor):
        x=self.block1(x)
        x=self.block2(x)
        conv3_1=self.conv3_1(x)
        x=self.block3(conv3_1)
        conv4_1=self.conv4_1(x)
        x=self.block4(conv4_1)
        conv5_1=self.conv5_1(x)
        out=self.block5(conv5_1)

        return conv3_1,conv4_1,conv5_1,out

class VggFlattener(nn.Module):
    def __init__(self,parm:Parameter):

        super(VggFlattener, self).__init__()
        self.parm=parm
        self.flatten=nn.Sequential(
            nn.Conv2d(512,256,3,2,0),
            nn.InstanceNorm2d(256),
            nn.LeakyReLU(0.2,True),
            nn.Conv2d(256,256,3,1,0),
            nn.InstanceNorm2d(256),
            nn.LeakyReLU(0.2,True)
        )
    def forward(self, x):
        return self.flatten(x).view(-1,self.parm.zdim)


#######  解码器体系结构    ###############
class Decoders(nn.Module):
    def __init__(self,parm:Parameter):
        super(Decoders, self).__init__()
        self.parm=parm
        self.idim=parm.idim
        self.wdim=parm.wdim
        self.decoderI=Decoder(parm,in_dim=parm.idim,out_channel=parm.channel,lb=0,ub=1)
        self.decoderW=Decoder(parm,in_dim=parm.wdim,out_channel=parm.channel,lb=0,ub=0.1)

        self.warper=Warper(parm)
        self.integrator=GridSpatialIntegral(parm)
        self.cutter=nn.Hardtanh(-1,1)

    def forward(self, zI:torch.Tensor,zW:torch.Tensor,baseGrid:torch.Tensor):
        self.texture=self.decoderI(zI.view(-1,self.parm.idim,1,1))
        self.diff_warp=self.decoderW(zW.view(-1,self.wdim,1,1))

        self.warping=self.integrator(self.diff_warp) - 1.2
        self.warping=self.cutter(self.warping)
        self.resWarping=self.warping-baseGrid
        self.output=self.warper(self.texture,self.warping)

        return self.texture,self.resWarping,self.output,self.warping

class Decoder(nn.Module):
    def __init__(self,parm:Parameter,in_dim:int,out_channel:int,lb:float=0,ub:float=1):
        super(Decoder, self).__init__()
        self.parm=parm
        self.net=nn.Sequential(
            nn.ConvTranspose2d(in_channels=in_dim,out_channels=256,kernel_size=3,stride=1,padding=0),
            nn.InstanceNorm2d(256),
            nn.LeakyReLU(0.2,True),

            nn.ConvTranspose2d(in_channels=256,out_channels=512,kernel_size=3,stride=2,padding=0),
            nn.InstanceNorm2d(512),
            nn.LeakyReLU(0.2,True),

            # conv-block-5
            nn.ConvTranspose2d(512,512,kernel_size=2,stride=2,padding=0),
            nn.InstanceNorm2d(512),
            nn.LeakyReLU(0.2,True),
            nn.ConvTranspose2d(512,512,3,stride=1,padding=1),
            nn.InstanceNorm2d(512),
            nn.LeakyReLU(0.2,True),
            nn.ConvTranspose2d(512, 512, 3, stride=1, padding=1),
            nn.InstanceNorm2d(512),
            nn.LeakyReLU(0.2, True),
            nn.ConvTranspose2d(512, 512, 3, stride=1, padding=1),
            nn.InstanceNorm2d(512),
            nn.LeakyReLU(0.2, True),
            nn.ConvTranspose2d(512, 512, 3, stride=1, padding=1),
            nn.InstanceNorm2d(512),
            nn.LeakyReLU(0.2, True),

            # conv-block-4
            nn.ConvTranspose2d(512, 512, kernel_size=2, stride=2, padding=0),
            nn.InstanceNorm2d(512),
            nn.LeakyReLU(0.2, True),
            nn.ConvTranspose2d(512, 512, 3, stride=1, padding=1),
            nn.InstanceNorm2d(512),
            nn.LeakyReLU(0.2, True),
            nn.ConvTranspose2d(512, 512, 3, stride=1, padding=1),
            nn.InstanceNorm2d(512),
            nn.LeakyReLU(0.2, True),
            nn.ConvTranspose2d(512, 512, 3, stride=1, padding=1),
            nn.InstanceNorm2d(512),
            nn.LeakyReLU(0.2, True),
            nn.ConvTranspose2d(512, 256, 3, stride=1, padding=1),
            nn.InstanceNorm2d(256),
            nn.LeakyReLU(0.2, True),

            # conv-block-3
            nn.ConvTranspose2d(256, 256, kernel_size=2, stride=2, padding=0),
            nn.InstanceNorm2d(256),
            nn.LeakyReLU(0.2, True),
            nn.ConvTranspose2d(256, 256, 3, stride=1, padding=1),
            nn.InstanceNorm2d(256),
            nn.LeakyReLU(0.2, True),
            nn.ConvTranspose2d(256, 256, 3, stride=1, padding=1),
            nn.InstanceNorm2d(256),
            nn.LeakyReLU(0.2, True),
            nn.ConvTranspose2d(256, 256, 3, stride=1, padding=1),
            nn.InstanceNorm2d(256),
            nn.LeakyReLU(0.2, True),
            nn.ConvTranspose2d(256, 128, 3, stride=1, padding=1),
            nn.InstanceNorm2d(128),
            nn.LeakyReLU(0.2, True),

            # conv-block-2
            nn.ConvTranspose2d(128, 128, kernel_size=2, stride=2, padding=0),
            nn.InstanceNorm2d(128),
            nn.LeakyReLU(0.2, True),
            nn.ConvTranspose2d(128, 128, 3, stride=1, padding=1),
            nn.InstanceNorm2d(128),
            nn.LeakyReLU(0.2, True),
            nn.ConvTranspose2d(128, 64, 3, stride=1, padding=1),
            nn.InstanceNorm2d(64),
            nn.LeakyReLU(0.2, True),

            # conv-block-1
            nn.ConvTranspose2d(64, 64, kernel_size=2, stride=2, padding=0),
            nn.InstanceNorm2d(64),
            nn.LeakyReLU(0.2, True),
            nn.ConvTranspose2d(64, 64, 3, stride=1, padding=1),
            nn.InstanceNorm2d(64),
            nn.LeakyReLU(0.2, True),
            nn.ConvTranspose2d(64, out_channel, 3, stride=1, padding=1),
            nn.InstanceNorm2d(3),
            nn.Hardtanh(lb,ub)
        )

    def forward(self, x):
        return self.net(x)